# -*- coding: utf-8  -*-
# @Time : 2021/8/30  20:25
# @Author : zhangnengbo
# @File : jindi_830gai_dataset.py
# @Company : HPY
# 美剧字幕
# http://www.tvsubtitles.net/subtitle-219-1-en.html

path = 'D:\\data\\jindi\\SAVE'

import os
import sys
import pickle

from skimage import io
import matplotlib.pyplot as plt
import numpy
import torch
from torch.utils.data import Dataset
import random
from PIL import Image
import cv2
import numpy as np



class StoneTrain(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        self.path = path
        self.train = []
        with open('G:\\data\\jindi\\SAVE\\JinDiTrain.txt', 'r') as f:
            while True:
                content = f.readline()
                if not content:
                    break
                self.train.append(content)

        random.shuffle(self.train)
        self.transform = transform

    def __len__(self):
        return len(self.train)

    def __getitem__(self, index):
        # label = self.data['fine_labels'.encode()][index]
        # r = self.data['data'.encode()][index, :1024].reshape(32, 32)
        # g = self.data['data'.encode()][index, 1024:2048].reshape(32, 32)
        # b = self.data['data'.encode()][index, 2048:].reshape(32, 32)
        content_ori = self.train[index]
        content = content_ori.split(' ')
        label_content = content[-1].strip('/n')
        # print(content)
      #  print(content_ori[0:-3])
        label = int(label_content)
        fileName = content_ori[0:-3]
        image = np.array(cv2.imread(fileName))
        image = image[:,:,1]
        # try:
        # image = image[:, :, ::-1]
        #except:
        #    print(content_ori)
        image = cv2.merge([image, image, image])
        image = Image.fromarray(np.uint8(image))
        image = image.resize((96, 96))

        if self.transform:
            image = self.transform(image)
        image= image[1,:,:]
        image = image.unsqueeze(0)
        return image, label

class StoneTest(Dataset):
    """cifar100 test dataset, derived from
    torch.utils.data.DataSet
    """

    def __init__(self, path, transform=None):
        self.path = path
        self.test = []
        with open('G:\\data\\jindi\\SAVE\\JinDiTest.txt', 'r') as f:
            while True:
                content = f.readline()
                if not content:
                    break
                self.test.append(content)

        random.shuffle(self.test)
        self.transform = transform

    def __len__(self):
        return len(self.test)

    def __getitem__(self, index):
        # label = self.data['fine_labels'.encode()][index]
        # r = self.data['data'.encode()][index, :1024].reshape(32, 32)
        # g = self.data['data'.encode()][index, 1024:2048].reshape(32, 32)
        # b = self.data['data'.encode()][index, 2048:].reshape(32, 32)
        content_ori = self.test[index]
        content = content_ori.split(' ')
        label_content = content[-1].strip('/n')
       # print(content[0:-3])
        label = int(label_content)
        image = np.array(cv2.imread(content_ori[0:-3]))
        image = image[:, :, 1]
        image = cv2.merge([image, image, image])
        # image = image[:, :, ::-1]

        image = Image.fromarray(np.uint8(image))
        image = image.resize((96, 96))

        if self.transform:
            image = self.transform(image)
        image = image[1, :, :]
        image = image.unsqueeze(0)
        return image, label
